import qualified Data.Map as Map

-- 0, 3, 6, 0, 3, 3, 1, 0, 4, 0, ...

type Index = Int
type Value = Int

type IndexRecord = Either Index (Index, Index)
type Memory = Map.Map Value IndexRecord
data State = State Value Memory deriving (Show)

updateMemory :: Memory -> (Index,Value) -> Memory
updateMemory memory (index, value) = 
    Map.insert value indexRecord memory
    where
        indexRecord = case Map.lookup value memory of 
            Nothing -> (Left index)
            (Just (Left idx)) -> (Right (index, idx))
            (Just (Right (lastIdx, llastIdx))) -> (Right (index, lastIdx))

unwrap :: Maybe a -> a
unwrap (Just v) = v

step :: State -> Index -> State
step (State lastVal memory) idx = 
    State insertValue (updateMemory memory (idx, insertValue))
    where
        indexRecord = unwrap $ Map.lookup lastVal memory        
        insertValue = case indexRecord of
            (Left _) -> 0
            (Right (lastIdx, llastIdx)) -> lastIdx - llastIdx

memoryGame :: [Value] -> State
memoryGame numbers = 
    foldl step initState [(length numbers)..(30000000-1)]
    where
        initMemory = foldl (\mem (i,v) -> Map.insert v (Left i) mem) Map.empty . zip [0..] $ numbers
        initState = State (last numbers) initMemory

getValue :: State -> Value
getValue (State value _) = value

main :: IO ()
main = do
    print . getValue . memoryGame $ [18,8,0,5,4,1,20]